In [1]:

    

GPU_NUMBER = 3
root = '/data/vision/torralba/health-habits/other/enes/'

%matplotlib inline
import matplotlib.pyplot as plt

import os
import sys
import random
import json
import math
import scipy.misc
import ujson, gzip
sys.path.append( root + 'Utils/')

import pandas as pd
import numpy as np
import tensorflow as tf

from PIL import Image
from IPython.display import display
from pprint import pprint
from notebook_utils import *
from skimage import color, io


os.environ['CUDA_VISIBLE_DEVICES'] = str(GPU_NUMBER);
data_root = '/data/vision/torralba/yusuf/imagenet/data/images/train256/'

import warnings
warnings.filterwarnings('ignore')


    

In [2]:

    

f = gzip.open('/afs/csail.mit.edu/u/k/kocabey/TF_colorization/dataset_indexes/imagenet_train_256_saturation_values.json.gz', 'rt')
saturation_index = ujson.load(f)
len(saturation_index)
all_paths = [k for (k,v) in saturation_index.items() if v > 0.1]
print len(all_paths)


    

    




1178311

In [3]:

    

quantized_array = np.load('/afs/csail.mit.edu/u/k/kocabey/TF_colorization/network/pts_in_hull.npy')


    

In [4]:

    

def gaussian( x, var ):
  return np.exp( -(x**2) / (2 * var**2))


    

In [5]:

    

def get_data(path):
  img = io.imread(path)
  resized_img = scipy.misc.imresize(img, (64,64))
  img = color.rgb2lab(img)
  assert img.shape == (256,256,3)
  image = img[:,:,0:1]
  img = color.rgb2lab(resized_img)
  colors = img[:,:,1:3]
  colors = np.tile( colors.reshape((64,64,1,2)), (1,1,313,1))
  big_quantized = np.tile( quantized_array, (64,64,1,1))
  distances = np.linalg.norm(colors - big_quantized, axis = 3)
  d = distances.copy()
  d.sort(axis = 2)
  low_values = (distances > np.tile( d[:,:,4:5], (1,1,313) ))
  gaussian_distances = gaussian(distances, 5)
  gaussian_distances[low_values] = 0
  output = gaussian_distances / np.sum(gaussian_distances, axis = 2).reshape((64,64,1))
  return image, output

image, output = get_data( data_root + all_paths[0] )
# print output[0][0]
Image.fromarray( image.reshape((256,256)).astype(np.uint8))


    

    
Out[5]:

In [6]:

    

import numpy as np
import scipy.misc

from skimage import color, io

ab_to_dist = {}


def image_path_to_image_and_distribution_tensor(path):
    '''Converts an image path to a LAB image and a [64, 64, 313] tensor of color distribution values.'''
    raw_img = io.imread(path)
    lab_img = color.rgb2lab(raw_img)

    img = lab_img[:, :, 1:3]
    dist = np.zeros([64, 64, 313])

    h, w, _ = dist.shape
    for x in xrange(w):
        for y in xrange(h):
            dist[x][y] = _map_ab_to_distribution(np.floor(img[4 * x][4 * y]).tolist())

    return lab_img[:, :, 0], dist


def _gaussian(x, var):
    '''Gaussian on np array'''
    return np.exp(-(x**2) / (2 * var**2))


def _precompute_distributions():
    '''Precomputes the distribution we want for each integer a,b value.'''
    global ab_to_dist
    print "Precomputing distributions... will take a second"

    quantized_array = np.load('/afs/csail.mit.edu/u/k/kocabey/TF_colorization/network/pts_in_hull.npy')

    for a in xrange(-120, 120):
        for b in xrange(-120, 120):
            tiled_ab = np.tile([a, b], (313, 1))

            distances = np.linalg.norm(quantized_array - tiled_ab, axis=1)
            d = distances.copy()
            d.sort()

            low_values = (distances > np.tile(d[4], (313)))
            gaussian_distances = _gaussian(distances, 5)
            gaussian_distances[low_values] = 0

            dist = gaussian_distances / np.sum(gaussian_distances)

            ab_to_dist[(a, b)] = dist
    print "Done"


def _map_ab_to_distribution(ab):
    '''Map an integer (a,b) tuple to a 313 deep distribution.'''
    if len(ab_to_dist) == 0:
        _precompute_distributions()

    return ab_to_dist[ab]


    

In [7]:

    

def get_colorized_image( image, prediction ):
  T = 0.38
  epsilon = 1e-8

  annealed_mean = np.exp( np.log(prediction + epsilon) / T ) 
  annealed_mean /= np.sum(annealed_mean, axis = 2).reshape((256,256,1))
  
  predicted_coloring = np.dot(annealed_mean, quantized_array)
  colorized_image = np.zeros( (256,256,3) )
  colorized_image[:,:,0:1] = image
  colorized_image[:,:,1:] = predicted_coloring
  
  return Image.fromarray( (255 * color.lab2rgb(colorized_image)).astype(np.uint8) )

def upscale(distribution):
  output_ = tf.placeholder(tf.float32, shape=[64, 64, 313])
  upscaled = tf.image.resize_images(output_, [256, 256])
  with tf.Session() as sess:
    return sess.run(upscaled, feed_dict={output_: distribution})

%timeit image, output = get_data( data_root + all_paths[0] )
%timeit image2, output2 = image_path_to_image_and_distribution_tensor( data_root + all_paths[0] )

display(get_colorized_image(image, upscale(output)))
display(get_colorized_image(image2.reshape((256, 256, 1)), upscale(output2)))


    

    




10 loops, best of 3: 152 ms per loop
Precomputing distributions... will take a second






    




---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-7-d99b3a6cb11f> in <module>()
     20 
     21 get_ipython().magic(u'timeit image, output = get_data( data_root + all_paths[0] )')
---> 22 get_ipython().magic(u'timeit image2, output2 = image_path_to_image_and_distribution_tensor( data_root + all_paths[0] )')
     23 
     24 display(get_colorized_image(image, upscale(output)))

/afs/csail.mit.edu/u/k/kocabey/.virtualenvs/tensorflow/local/lib/python2.7/site-packages/IPython/core/interactiveshell.pyc in magic(self, arg_s)
   2156         magic_name, _, magic_arg_s = arg_s.partition(' ')
   2157         magic_name = magic_name.lstrip(prefilter.ESC_MAGIC)
-> 2158         return self.run_line_magic(magic_name, magic_arg_s)
   2159 
   2160     #-------------------------------------------------------------------------

/afs/csail.mit.edu/u/k/kocabey/.virtualenvs/tensorflow/local/lib/python2.7/site-packages/IPython/core/interactiveshell.pyc in run_line_magic(self, magic_name, line)
   2077                 kwargs['local_ns'] = sys._getframe(stack_depth).f_locals
   2078             with self.builtin_trap:
-> 2079                 result = fn(*args,**kwargs)
   2080             return result
   2081 

<decorator-gen-58> in timeit(self, line, cell)

/afs/csail.mit.edu/u/k/kocabey/.virtualenvs/tensorflow/local/lib/python2.7/site-packages/IPython/core/magic.pyc in <lambda>(f, *a, **k)
    186     # but it's overkill for just that one bit of state.
    187     def magic_deco(arg):
--> 188         call = lambda f, *a, **k: f(*a, **k)
    189 
    190         if callable(arg):

/afs/csail.mit.edu/u/k/kocabey/.virtualenvs/tensorflow/local/lib/python2.7/site-packages/IPython/core/magics/execution.pyc in timeit(self, line, cell)
   1042             number = 1
   1043             for _ in range(1, 10):
-> 1044                 time_number = timer.timeit(number)
   1045                 worst_tuning = max(worst_tuning, time_number / number)
   1046                 if time_number >= 0.2:

/afs/csail.mit.edu/u/k/kocabey/.virtualenvs/tensorflow/local/lib/python2.7/site-packages/IPython/core/magics/execution.pyc in timeit(self, number)
    137         gc.disable()
    138         try:
--> 139             timing = self.inner(it, self.timer)
    140         finally:
    141             if gcold:

<magic-timeit> in inner(_it, _timer)

<ipython-input-6-f9484c356161> in image_path_to_image_and_distribution_tensor(path)
     18     for x in xrange(w):
     19         for y in xrange(h):
---> 20             dist[x][y] = _map_ab_to_distribution(np.floor(img[4 * x][4 * y]).tolist())
     21 
     22     return lab_img[:, :, 0], dist

<ipython-input-6-f9484c356161> in _map_ab_to_distribution(ab)
     56     '''Map an integer (a,b) tuple to a 313 deep distribution.'''
     57     if len(ab_to_dist) == 0:
---> 58         _precompute_distributions()
     59 
     60     return ab_to_dist[ab]

<ipython-input-6-f9484c356161> in _precompute_distributions()
     49             dist = gaussian_distances / np.sum(gaussian_distances)
     50 
---> 51             ab_to_dist[[a, b]] = dist
     52     print "Done"
     53 

TypeError: unhashable type: 'list'

In [ ]: